import jsonpath


class Extractor(object):
    """
    提取器
    """

    def __init__(self):
        # 初始化声明一个字典用于存储变量
        self.variable_pool = {}

    def extract_output(self, res=None):
        """
        提取返回值参数
        """
        for key, value in self.variable_pool.items():
            extract_value = self.extract_jsonpath(res, value)
            return {key: extract_value}

    def format_output(self, output_variables):
        """
        格式化参数
        """
        if isinstance(output_variables, list):
            for item in output_variables:
                self.format_output(item)
        elif isinstance(output_variables, dict):
            for key, value in output_variables.items():
                self.variable_pool.update({key: value})
        elif isinstance(output_variables, str):
            self.variable_pool.update({output_variables: output_variables})
        else:
            raise Exception('未知headers参数格式')

    @staticmethod
    def extract_jsonpath(res_obj=None, expr=None):
        result = jsonpath.jsonpath(res_obj, expr)
        if result == False:
            print('未提取到值')
        elif isinstance(result, list):
            if len(result) == 1:
                result = result[0]
        return result
